package cn.doitedu.dw_etl.entropy

import org.apache.spark.sql.Row
import org.apache.spark.sql.expressions.{MutableAggregationBuffer, UserDefinedAggregateFunction}
import org.apache.spark.sql.types.{DataType, DataTypes, StructType}

object EntropyUDAF extends UserDefinedAggregateFunction{
  // 定义函数接收的参数：参数有几个，以及参数的字段类型
  override def inputSchema: StructType = {
    new StructType()
      .add("pay_cat",DataTypes.StringType)
      .add("pay_rate",DataTypes.DoubleType)
  }

  // 中间结果的缓存数据的存储结构
  override def bufferSchema: StructType = {
    new StructType()
      .add("mp",DataTypes.createMapType(DataTypes.StringType,DataTypes.DoubleType))
  }

  // 函数的最终返回值类型
  override def dataType: DataType = DataTypes.DoubleType

  // 计算逻辑针对不同数据是否每次都能得到不同结果
  override def deterministic: Boolean = true

  // 初始化缓存数据
  override def initialize(buffer: MutableAggregationBuffer): Unit = {
    buffer.update(0,Map.empty[String,Double])
  }

  // 每进来一行数据，如何更新缓存（核心计算逻辑所在）
  override def update(buffer: MutableAggregationBuffer, input: Row): Unit = {
    // 从输入行中取出所需的两个字段值
    val pay_cat: String = input.getAs[String](0)
    val pay_rate: Double = input.getAs[Double](1)

    // 从缓存中取出hashmap
    var mp: collection.Map[String, Double] = buffer.getMap[String, Double](0)

    // 更新缓存hashmap
    val oldV = mp.getOrElse(pay_cat,0.0)
    val newV = oldV + pay_rate
    mp += ((pay_cat,newV))

    // 把更新后的map，重新放入缓存
    buffer.update(0,mp)
  }

  // 分区之间的全局聚合
  override def merge(buffer1: MutableAggregationBuffer, buffer2: Row): Unit = {

    var mp1: collection.Map[String, Double] = buffer1.getMap[String, Double](0)
    var mp2: collection.Map[String, Double] = buffer2.getMap[String, Double](0)

    // 合并两个map
    for ((k:String,v:Double) <- mp2) {
      val newV: Double = mp1.getOrElse(k,0.0) + v
      mp1 += ((k,newV))
    }

    // 更新buffer1
    buffer1.update(0,mp1)

  }

  // 输出最终结果
  override def evaluate(buffer: Row): Any = {
    // buffer:
    //  男->sum1,  女->sum2

    val mp = buffer.getMap[String, Double](0)

    // 赔付率总和
    val total = mp.values.sum

    var entro = 0.0
    for ((pay_cat,pay_rate_sum) <- mp) {
      val fenmu = Math.log(pay_rate_sum) * (pay_rate_sum/total) * Math.log((pay_rate_sum/total))
      if(fenmu != 0) {
        entro = entro + (-1) / fenmu
      }
    }
    entro
  }
}
